import torch
from torch.utils.data import Dataset

# from torch.utils.data import DataLoader
import numpy as np
import pandas


class ExcelDataset(Dataset):
    def __init__(self, filepath="data.xlsx", sheet_name=0):

        print(f"reading {filepath}, sheet={sheet_name}")

        df = pandas.read_excel(
            filepath,
            header=0,
            index_col=0,
            names=["feat1", "feat2", "label"],
            sheet_name=sheet_name,
            dtype={"feat1": np.float32, "feat2": np.float32, "label": np.int32},
        )

        print(f"the shape of dataframe is {df.shape}")

        feat = df.iloc[:, :2].values
        label = df.iloc[:, 2].values

        self.x = torch.from_numpy(feat)
        self.y = torch.from_numpy(label)

    def __len__(self):
        return len(self.y)

    def __getitem__(self, index):
        return self.x[index], self.y[index]


"""
Demo:
excel_dataset = ExcelDataset(sheet_name="sheet_name")
excel_dataloader = DataLoader(excel_dataset, batch_size=8, shuffle=True)
for idx, (batch_x, batch_y) in enumerate(excel_dataloader):
    pass
"""


class CsvDataset(Dataset):
    def __init__(self, filepath="data.csv"):

        print(f"reading {filepath}")

        df = pandas.read_csv(
            filepath,
            header=0,
            index_col=0,  # type: ignore
            encoding="utf-8",
            names=["feat1", "feat2", "label"],  # type: ignore
            dtype={"feat1": np.float32, "feat2": np.float32, "label": np.int32},
            skip_blank_lines=True,
        )
        print(f"the shape of dataframe is {df.shape}")  # type: ignore

        feat = df.iloc[:, :2].values  # type: ignore
        label = df.iloc[:, 2].values  # type: ignore

        self.x = torch.from_numpy(feat)
        self.y = torch.from_numpy(label)

    def __len__(self):
        return len(self.y)

    def __getitem__(self, index):
        return self.x[index], self.y[index]


class Csv2Dataset(Dataset):
    def __init__(self, filepath="data.csv"):

        print(f"reading {filepath}")

        with open(filepath, encoding="utf-8") as f:
            lines = f.readlines()

        feat = []
        label = []
        for line in lines[1:]:
            values = line.strip().split(",")
            row_feat = [float(v) for v in values[1:3]]
            row_label = int(values[3])

            feat.append(row_feat)
            label.append(row_label)

        feat = np.array(feat, dtype=np.float32)
        label = np.array(label, dtype=np.int32)

        self.x = torch.from_numpy(feat)
        self.y = torch.from_numpy(label)

    def __len__(self):
        return len(self.y)

    def __getitem__(self, index):
        return self.x[index], self.y[index]
